import datetime
import json
import sys
from absl import flags

from selenium import webdriver

from coin.exchange.base.kr_rest.auth_base import AuthKey

FLAGS = flags.FLAGS

flags.DEFINE_string("owner", None, "owner")

flags.DEFINE_string("key_file", None, "key_file")


class LocalStorage(object):
  def __init__(self, driver):
    assert type(driver) == webdriver.Chrome
    self.driver = driver

  def set(self, key, value):
    self.driver.execute_script(\
      "window.localStorage.setItem('{}',{})".format(key, json.dumps(value)))

  def get(self, key=None):
    if key:
      return self.driver.execute_script(\
        "return window.localStorage.getItem('{}')".format(key))
    else:
      return self.driver.execute_script("""
        var items = {}, ls = window.localStorage;
        for (var i = 0, k; i < ls.length; i++)
          items[k = ls.key(i)] = ls.getItem(k);
        return items;
        """)

  def remove(self, key):
    self.driver.execute_script(\
      "window.localStorage.removeItem('{}');".format(key))

  def clear(self):
    self.driver.execute_script(\
      "window.localStorage.clear();")


def _create_driver():
  driver = webdriver.Chrome('chromedriver')
  driver.implicitly_wait(3)
  driver.get('https://upbit.com')
  driver.implicitly_wait(8)
  driver.get("https://upbit.com/signin")
  return driver


class UpbitChromeDriver(object):
  def __init__(self):
    self.driver = _create_driver()
    self.refresh_token = None

  def get_key(self, owner=None):
    storage = LocalStorage(self.driver)
    obj = storage.get("bitmochi")
    obj = json.loads(obj)
    key = AuthKey(owner=owner,
                  key_name="view+trade",
                  access_key=obj["sessions"]["tokens"]["access_key"],
                  secret_key=obj["sessions"]["tokens"]["secret_key"],
                  raw=dict(refresh_token=obj["sessions"]["session"]["refreshToken"]),
                  generated=str(datetime.datetime.utcnow()),
                  generated_human=str(datetime.datetime.now()))
    return key


def main():
  driver = UpbitChromeDriver()
  input("Press Enter to continue...")
  key = driver.get_key(FLAGS.owner)

  print(key.as_json_str())
  if FLAGS.key_file is not None:
    for key_file in FLAGS.key_file.split(","):
      try:
        print(key_file)
        key.save_file(key_file)
      except Exception:
        print("failed to save to", key_file)
  input("Press Enter to continue...")


if __name__ == '__main__':
  FLAGS(sys.argv)
  main()
